'''
K-band algorithm for processing sequences.
Author: wym6912
Time: 2020-10-02 15:00
'''
import argparse

Protein = "CSTPAGNDEQHRKMILVFYW?"
DNA = "ATCGN?"

#use BLOSUM62
protein_seq = [[9,-1,-1,-3, 0,-3,-3,-3,-4,-3,-3,-3,-3,-1,-1,-1,-1,-2,-2,-2],[-1, 4, 1,-1, 1, 0, 1, 0, 0, 0,-1,-1, 0,-1,-2,-2,-2,-2,-2,-3],[-1, 1, 5,-1, 0,-2, 0,-1,-1,-1,-2,-1,-1,-1,-1,-1, 0,-2,-2,-2],[-3,-1,-1, 7,-1,-2,-2,-1,-1,-1,-2,-2,-1,-2,-3,-3,-2,-4,-3,-4],[ 0, 1, 0,-1, 4, 0,-2,-2,-1,-1,-2,-1,-1,-1,-1,-1, 0,-2,-2,-3],[-3, 0,-2,-2, 0, 6, 0,-1,-2,-2,-2,-2,-2,-3,-4,-4,-3,-3,-3,-2],[-3, 1, 0,-2,-2, 0, 6, 1, 0, 0, 1, 0, 0,-2,-3,-3,-3,-3,-2,-4],[-3, 0,-1,-1,-2,-1, 1, 6, 2, 0,-1,-2,-1,-3,-3,-4,-3,-3,-3,-4],[-4, 0,-1,-1,-1,-2, 0, 2, 5, 2, 0, 0, 1,-2,-3,-3,-2,-3,-2,-3],[-3, 0,-1,-1,-1,-2, 0, 0, 2, 5, 0, 1, 1, 0,-3,-2,-2,-3,-1,-2],[-3,-1,-2,-2,-2,-2, 1,-1, 0, 0, 8, 0,-1,-2,-3,-3,-3,-1, 2,-2],[-3,-1,-1,-2,-1,-2, 0,-2, 0, 1, 0, 5, 2,-1,-3,-2,-3,-3,-2,-3],[-3, 0,-1,-1,-1,-2, 0,-1, 1, 1,-1, 2, 5,-1,-3,-2,-2,-3,-2,-3],[-1,-1,-1,-2,-1,-3,-2,-3,-2, 0,-2,-1,-1, 5, 1, 2, 1, 0,-1,-1],[-1,-2,-1,-3,-1,-4,-3,-3,-3,-3,-3,-3,-3, 1, 4, 2, 3, 0,-1,-3],[-1,-2,-1,-3,-1,-4,-3,-4,-3,-2,-3,-2,-2, 2, 2, 4, 1, 0,-1,-2],[-1,-2, 0,-2, 0,-3,-3,-3,-2,-2,-3,-3,-2, 1, 3, 1, 4,-1,-1,-3],[-2,-2,-2,-4,-2,-3,-3,-3,-3,-3,-1,-3,-3, 0, 0, 0,-1, 6, 3, 1],[-2,-2,-2,-3,-2,-3,-2,-3,-2,-1, 2,-2,-2,-1,-1,-1,-1, 3, 7, 2],[-2,-3,-2,-4,-3,-2,-4,-4,-3,-2,-2,-3,-3,-1,-3,-2,-3, 1, 2,11]]

def parse_fasta(file_buffer):
    '''
    ignore the description line and return the sequence line, the file must have only 2 lines\n
    returns: string
    exception: TypeError: Check the file is illegal
    '''
    cnt = 0
    for data in file_buffer:
        if cnt == 1: return '?' + data.strip('\r\n')
        cnt += 1
    # If the file only have 1 line, then raise an Error
    raise TypeError("File is illegal!")

def parse_file():
    '''
    define a parser to process file name\n
    returns: processed arguments
    '''
    parser = argparse.ArgumentParser(description = "K-band algorithm parser. We can process fasta file.")
    parser.add_argument("--file1", type = str, help = "K-band file 1")
    parser.add_argument("--file2", type = str, help = "K-band file 2")
    parser.add_argument("-i", "--ignore-letter-case", type = bool, default = False, help = "Treat the upper letters as lower letters")
    parser.add_argument("-t", "--type", type = int, default = 1, help = "Processing data type: 1=DNA, 2=Protein(Use BLOSUM62 Matrix)")
    parser.add_argument("-d", type = int, default = 2, help = "Character aligns to a gap penalty (>0)")
    parser.add_argument("-g", type = int, default = 1, help = "Expand aligns penalty (>0)")
    parser.add_argument("-s", "--success", type = int, default = 4, help = "Success aligns score")
    parser.add_argument("-b", "--band", type = int, default = 2, help = "the arugments of k-band")
    return parser.parse_args()

def check_seq(s1, s2, type):
    '''
    check the length of `s1` is equal to `s2`\n
    input: two sequences \n
    return: None \n
    expection: TypeError, the data mismatch the data rule(DNA only 5 cases, protein 20 cases);\n
    ..       : TypeError, the type not in [1, 2]
    '''
    if type == 1: # Type DNA
        for i in s1:
            if i.upper() not in DNA:
                raise TypeError("Sequence 1 mismatches DNA rule, the ASCII of unmatched char {0:} is {1:}".format(i, ord(i)))
        for i in s2:
            if i.upper() not in DNA:
                raise TypeError("Sequence 2 mismatches DNA rule, the ASCII of unmatched char {0:} is {1:}".format(i, ord(i)))
    elif type == 2: # Type protein
        for i in s1:
            if i.upper() not in Protein:
                raise TypeError("Sequence 1 mismatches protein rule")
        for i in s2:
            if i.upper() not in Protein:
                raise TypeError("Sequence 2 mismatches protein rule")
    else: raise TypeError("The type is not in [1, 2]")

def mapping(x, y, n):
    '''
    mapping coordinate into a number\n
    input: x, y (on matrix), n (the size of matrix)\n
    return: number in matrix
    expection: TypeError, x or y out of [0, n]
    '''
    if not (0 <= x <= n and 0 <= y <= n): raise TypeError("Coordinate out of range [0, " + str(n) + "]")
    return x * n + y

def insideBand(i, j, band):
    '''
    judge the coordinate in the band
    input: i, j, band
    return: True or False
    '''
    return -band <= i - j <= band

def findPlace(c):
    '''
    find protein character in Protein sequence
    input: c
    return: place on Protein array
    exception: TypeError
    '''
    for i in range(20):
        if Protein[i] == c:
            return i
    raise TypeError("Protein out of range")

def same(c1, c2, seq_type, ignore_letter_case, succ, fail):
    '''
    calcuate the score of c1 and c2
    input: c1, c2, seq_type, ignore letter case
    return: score
    '''
    if ignore_letter_case: c1.upper(); c2.upper()
    if seq_type == 1: # DNA
        if c1 == c2: return succ
        else: return -fail
    else: # Protein
        if c1 == '?': return 0
        elif c2 == '?': return 0
        return protein_seq[findPlace(c1)][findPlace(c2)]

def align(s1, s2, ignore_letter_case, seq_type, length, d, g, band, succ):
    '''
    align two sequences using kband algorithm\n
    input: two sequences which the sequences are approximatly similar; ignore letter case (True or False); type (DNA or protein); 
    the length of string; d, g\n
    return: [alignment_score, alignment_sequence1, alignment_sequence2]
    '''
    if len(s1) != len(s2):
        if len(s2) > len(s1): s1, s2 = s2, s1
        s2 = s2 + (len(s1) - len(s2)) * '?'
    #print(len(s2), len(s1), length)
    matrix = {}
    # Initialize and align
    matrix[mapping(0, 0, length)] = 0
    for i in range(1, length + 1):
        matrix[mapping(i, 0, length)] = -i * d
        matrix[mapping(0, i, length)] = -i * d
    score = 0
    for i in range(1, length + 1):
        for h in range(-band, band + 1):
            j = i + h
            if 1 <= j <= length:
                matrix[mapping(i, j, length)] = matrix[mapping(i - 1, j - 1, length)] + same(s1[i], s2[j], seq_type, ignore_letter_case, succ, d)
                if insideBand(i - 1, j, band): 
                    matrix[mapping(i, j, length)] = max(matrix[mapping(i - 1, j, length)] - d, matrix[mapping(i, j, length)])
                if insideBand(i, j - 1, band): 
                    matrix[mapping(i, j, length)] = max(matrix[mapping(i, j - 1, length)] - d, matrix[mapping(i, j, length)])
                score = max(score, matrix[mapping(i, j, length)])
                #print("Score[i][j]", i, j, matrix[mapping(i, j, length)])
    # Traceback
    L1 = ''
    L2 = ''
    i, j = length, length
    while(i != 0 and j != 0):
        if s1[i] == '?':
            i = i - 1
            continue
        elif s2[j] == '?':
            j = j - 1
            continue
        if insideBand(i - 1, j, band):
            if matrix[mapping(i, j, length)] == matrix[mapping(i - 1, j, length)] - d:
                L1 = L1 + s1[i]
                L2 = L2 + "-"
                i = i - 1
                continue
        if insideBand(i, j - 1, band):
            if matrix[mapping(i, j, length)] == matrix[mapping(i, j - 1, length)] - d:
                L1 = L1 + "-"
                L2 = L2 + s2[j]
                j = j - 1
                continue
        try:
            if(matrix[mapping(i, j, length)] == matrix[mapping(i - 1, j - 1, length)] + same(s1[i], s2[j], seq_type, ignore_letter_case, succ, d)):
                L1 = L1 + s1[i]
                L2 = L2 + s2[j]
                i = i - 1
                j = j - 1
        except KeyError: print(i, j)
    #print(i, j)
    while i != 0:
        L2 = L2 + '-'
        L1 = L1 + s1[i]
        i -= 1
    while j != 0:
        L1 = L1 + '-'
        L2 = L2 + s2[j]
        j -= 1
    L1 = L1[::-1]
    L2 = L2[::-1]
    return [score, L1, L2]

if __name__ == '__main__':
    args = parse_file()
    f1 = open(args.file1, "r")
    f2 = open(args.file2, "r")
    # Read File
    s1 = parse_fasta(f1)
    s2 = parse_fasta(f2)
    check_seq(s1, s2, args.type)
    # Process alignment
    # length of s is len(s) - 1 because of added a space before the data
    score, s1, s2 = align(s1.strip(), s2.strip(), args.ignore_letter_case, args.type, len(s1) - 1, args.d, args.g, args.band, args.success)
    #print(len(s1), len(s2))
    print("Score =", score, "\nAlignment s1 = \n", s1, "\nAlignment s2 = \n", s2)
        